import os.path
import random
import numpy as np
import cv2
import torch
import torch.utils.data as data
from . import utils


class LRHRDataset(data.Dataset):
    '''read HR,LR images for train'''

    def __init__(self, mode, exp):
        super(LRHRDataset, self).__init__()
        self.paths_LR = None
        self.paths_HR = None
        if mode == "train":
            dataroot_HR = "./datasets/train/"+exp+"/HR"
            dataroot_LR = "./datasets/train/"+exp+"/LR"
        elif mode == "val":
            dataroot_HR = "./datasets/test/"+exp+"/HR"
            dataroot_LR = "./datasets/test/"+exp+"/LR"

        self.paths_HR = utils.get_image_paths(dataroot_HR)
        self.paths_LR = utils.get_image_paths(dataroot_LR)

        if self.paths_HR and self.paths_LR:
            assert len(self.paths_LR) == len(self.paths_HR), \
                'HR and LR datasets have different number of images - {}, {}.'.format(
                    len(self.paths_LR), len(self.paths_HR))

        self.random_scale_list = [1]
        self.mode = mode

    def __getitem__(self, index):
        HR_path, LR_path = None, None
        scale = 4
        HR_size = 128

        HR_path = self.paths_HR[index]
        img_HR = utils.read_img(HR_path)

        LR_path = self.paths_LR[index]
        img_LR = utils.read_img(LR_path)

        # if the image size is too small
        H, W, _ = img_HR.shape
        if H < HR_size or W < HR_size:
            img_HR = cv2.resize(
                np.copy(img_HR), (HR_size, HR_size), interpolation=cv2.INTER_LINEAR)
            # using matlab imresize
            img_LR = utils.imresize_np(img_HR, 1 / scale, True)
            if img_LR.ndim == 2:
                img_LR = np.expand_dims(img_LR, axis=2)

        H, W, C = img_LR.shape
        LR_size = HR_size // scale
        if self.mode == "train":
            # randomly crop
            rnd_h = random.randint(0, max(0, H - LR_size))
            rnd_w = random.randint(0, max(0, W - LR_size))
            img_LR = img_LR[rnd_h:rnd_h + LR_size, rnd_w:rnd_w + LR_size, :]
            rnd_h_HR, rnd_w_HR = int(rnd_h * scale), int(rnd_w * scale)
            img_HR = img_HR[rnd_h_HR:rnd_h_HR +
                            HR_size, rnd_w_HR:rnd_w_HR + HR_size, :]

            # augmentation - flip, rotate
            img_LR, img_HR = utils.augment([img_LR, img_HR], True, True)

        # BGR to RGB, HWC to CHW, numpy to tensor
        if img_HR.shape[2] == 3:
            img_HR = img_HR[:, :, [2, 1, 0]]
            img_LR = img_LR[:, :, [2, 1, 0]]
        img_HR = torch.from_numpy(np.ascontiguousarray(
            np.transpose(img_HR, (2, 0, 1)))).float()
        img_LR = torch.from_numpy(np.ascontiguousarray(
            np.transpose(img_LR, (2, 0, 1)))).float()

        return img_HR, img_LR

    def __len__(self):
        return len(self.paths_HR)


class LRDataset(data.Dataset):
    '''Read LR images for test'''

    def __init__(self, test_set):
        super(LRDataset, self).__init__()
        self.paths_LR = None
        dataroot_LR = "./datasets/test/"+test_set+"/LR"

        self.paths_LR = utils.get_image_paths(dataroot_LR)
        assert self.paths_LR, 'Error: LR paths are empty.'

    def __getitem__(self, index):
        LR_path = None

        # get LR image
        LR_path = self.paths_LR[index]
        img_LR = utils.read_img(LR_path)
        H, W, C = img_LR.shape

        # BGR to RGB, HWC to CHW, numpy to tensor
        if img_LR.shape[2] == 3:
            img_LR = img_LR[:, :, [2, 1, 0]]
        img_LR = torch.from_numpy(np.ascontiguousarray(
            np.transpose(img_LR, (2, 0, 1)))).float()

        return img_LR

    def __len__(self):
        return len(self.paths_LR)
